﻿using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;
using Exceptionless.Core;
using Newtonsoft.Json;

namespace Exceptionless.Web.Hubs;

public class WebSocketConnectionManager : IDisposable
{
    private static readonly ArraySegment<byte> _keepAliveMessage = new(Encoding.ASCII.GetBytes("{}"), 0, 2);
    private readonly ConcurrentDictionary<string, WebSocket> _connections = new();
    private readonly Timer? _timer;
    private readonly JsonSerializerSettings _serializerSettings;
    private readonly ILogger _logger;

    public WebSocketConnectionManager(AppOptions options, JsonSerializerSettings serializerSettings, ILoggerFactory loggerFactory)
    {
        _serializerSettings = serializerSettings;
        _logger = loggerFactory.CreateLogger<WebSocketConnectionManager>();
        if (!options.EnableWebSockets)
            return;

        _timer = new Timer(KeepAlive, null, TimeSpan.FromSeconds(10), TimeSpan.FromSeconds(10));
    }

    private void KeepAlive(object? state)
    {
        if (_connections is { IsEmpty: true, Count: 0 })
            return;

        Task.Factory.StartNew(async () =>
        {
            var sockets = GetAll();
            var openSockets = sockets.Where(s => s.State == WebSocketState.Open).ToArray();
            _logger.LogTrace("Sending web socket keep alive to {OpenSocketsCount} open connections of {SocketCount} total connections", openSockets.Length, sockets.Count);

            foreach (var socket in openSockets)
            {
                try
                {
                    await socket.SendAsync(buffer: _keepAliveMessage,
                        messageType: WebSocketMessageType.Text,
                        endOfMessage: true,
                        cancellationToken: CancellationToken.None);
                }
                catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
                {
                    // NOTE: This will not remove it from the ConnectionMappings.
                    await RemoveWebSocketAsync(socket);
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Error sending keep alive socket message: {Message}", ex.Message);
                }
            }
        });
    }

    public WebSocket? GetWebSocketById(string connectionId)
    {
        return _connections.TryGetValue(connectionId, out var socket) ? socket : null;
    }

    public ICollection<WebSocket> GetAll()
    {
        return _connections.Values;
    }

    public string GetConnectionId(WebSocket socket)
    {
        return _connections.FirstOrDefault(p => p.Value == socket).Key;
    }

    public string AddWebSocket(WebSocket socket)
    {
        string connectionId = Guid.NewGuid().ToString("N");
        _connections.TryAdd(connectionId, socket);
        return connectionId;
    }

    private Task RemoveWebSocketAsync(WebSocket socket)
    {
        string id = GetConnectionId(socket);
        if (String.IsNullOrEmpty(id) || !_connections.TryRemove(id, out var _))
            return Task.CompletedTask;

        return CloseWebSocketAsync(socket);
    }

    public Task RemoveWebSocketAsync(string id)
    {
        if (!_connections.TryRemove(id, out var socket))
            return Task.CompletedTask;

        return CloseWebSocketAsync(socket);
    }

    private async Task CloseWebSocketAsync(WebSocket socket)
    {
        if (!CanSendWebSocketMessage(socket))
            return;

        try
        {
            await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closed by manager", CancellationToken.None);
        }
        catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
        {
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error closing web socket: {Message}", ex.Message);
        }
    }

    private Task SendMessageAsync(WebSocket socket, object message)
    {
        if (!CanSendWebSocketMessage(socket))
            return Task.CompletedTask;

        string serializedMessage = JsonConvert.SerializeObject(message, _serializerSettings);
        Task.Factory.StartNew(async () =>
        {
            if (!CanSendWebSocketMessage(socket))
                return;

            try
            {
                await socket.SendAsync(buffer: new ArraySegment<byte>(Encoding.ASCII.GetBytes(serializedMessage), 0, serializedMessage.Length),
                    messageType: WebSocketMessageType.Text,
                    endOfMessage: true,
                    cancellationToken: CancellationToken.None);
            }
            catch (WebSocketException ex) when (ex.WebSocketErrorCode == WebSocketError.ConnectionClosedPrematurely)
            {
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, "Error sending socket message: {Message}", ex.Message);
            }
        });

        return Task.CompletedTask;
    }

    public Task SendMessageAsync(string connectionId, object message)
    {
        var socket = GetWebSocketById(connectionId);
        return socket is not null ? SendMessageAsync(socket, message) : Task.CompletedTask;
    }

    public Task SendMessageAsync(IEnumerable<string> connectionIds, object message)
    {
        return Task.WhenAll(connectionIds.Select(id =>
        {
            var socket = GetWebSocketById(id);
            return socket is not null ? SendMessageAsync(socket, message) : Task.CompletedTask;
        }));
    }

    public async Task SendMessageToAllAsync(object message, bool throwOnError = true)
    {
        foreach (var socket in GetAll())
        {
            if (!CanSendWebSocketMessage(socket))
                continue;

            try
            {
                await SendMessageAsync(socket, message);
            }
            catch (Exception)
            {
                if (throwOnError)
                    throw;
            }
        }
    }

    private bool CanSendWebSocketMessage(WebSocket socket)
    {
        return socket.State != WebSocketState.Aborted && socket.State != WebSocketState.Closed && socket.State != WebSocketState.CloseSent;
    }

    public void Dispose()
    {
        _timer?.Dispose();
    }
}
